import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

# 加载并显示原始图像
# animal
# root_path = './quality/animal/'
# image_paths = ['0001.png', '0300.png', '0500.png']  # 替换为你的图像路径
# map_paths = ['0001-origin.png', '0300-origin.png', '0500-origin.png', '0001-STASNet.png', '0300-STASNet.png', '0500-STASNet.png', '0001-TMFI.png', '0300-TMFI.png', '0500-TMFI.png', '0001-Mine.png', '0300-Mine.png', '0500-Mine.png']

# landscape
root_path = './quality/ablation/'
image_paths = ['0001.png', '0250.png', '0500.png']
map_paths = ['0001-origin.png', '0250-origin.png', '0500-origin.png', '0001-full.png', '0250-full.png', '0500-full.png', '0001-fusion.png', '0250-fusion.png', '0500-fusion.png', '0001-memory.png', '0250-memory.png', '0500-memory.png', '0001-mul.png', '0250-mul.png', '0500-mul.png']
for i in range(len(map_paths)):
    image_path = root_path + image_paths[i % 3]
    original_image = Image.open(image_path).convert('RGB')
    plt.figure(figsize=(15, 5))
    plt.subplot(1, 3, 1)
    plt.imshow(original_image)
    plt.title('Original Image')

    # 转换图像以便于后续处理
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    tensor_image = transform(original_image)

    # 加载并显示显著性图
    salience_map_path = root_path + map_paths[i]  # 替换为你的显著性图路径
    salience_map_image = Image.open(salience_map_path).convert('L')  # 以灰度模式打开
    plt.subplot(1, 3, 2)
    plt.imshow(salience_map_image, cmap='gray')
    plt.title('Salience Map')

    # 转换显著性图为张量并归一化
    tensor_salience_map = transform(salience_map_image)  # 形状为 (1, H, W)
    tensor_salience_map = (tensor_salience_map - tensor_salience_map.min()) / (tensor_salience_map.max() - tensor_salience_map.min())

    # 将显著性图转换为二维numpy数组
    salience_np = tensor_salience_map.squeeze(0).numpy()  # 形状 (H, W)

    # 应用热力图颜色映射
    cmap = plt.get_cmap('jet')
    heatmap = cmap(salience_np)[:, :, :3]  # 转换为RGB，去掉alpha通道
    heatmap = (heatmap * 255).astype(np.uint8)  # 调整到0-255范围

    # 创建一个与原始图像相同大小的RGBA图像用于存储热力图
    heatmap_pil = Image.new("RGBA", original_image.size)
    heatmap_pil.putdata([(*tuple(p), int(a * 255)) for p, a in zip(heatmap.reshape(-1, 3), salience_np.flatten())])

    # 叠加热力图和原始图像
    blended_image = Image.alpha_composite(original_image.convert("RGBA"), heatmap_pil)

    # 显示融合后的图像
    plt.subplot(1, 3, 3)
    plt.imshow(blended_image)
    plt.title('Blended Image with Salience Map')
    plt.show()

    # 保存融合后的图像
    output_path = salience_map_path + '-blended' + '.png'  # 输出图像路径
    blended_image.save(output_path)
    print(f"Blended image saved to {output_path}")



